# @Time : 2021/8/5 8:46
# @Author : Li Kunlun
# @Description : 读取和存储

from mxnet import nd
from mxnet.gluon import nn

# 1、读写NDArray
x = nd.ones(3)
# 存在文件名同为x的文件里
nd.save('x', x)

x2 = nd.load('x')
"""
将数据从存储文件读回内存
    [
    [ 1.  1.  1.]
    <NDArray 3 @cpu(0)>]
"""
print(x2)

y = nd.zeros(4)
nd.save('xy', [x, y])
x2, y2 = nd.load('xy')

"""
存储一列NDArray并读回内存
    (
    [ 1.  1.  1.]
    <NDArray 3 @cpu(0)>, 
    [ 0.  0.  0.  0.]
    <NDArray 4 @cpu(0)>)
"""
print((x2, y2))

mydict = {'x': x, 'y': y}
nd.save('mydict', mydict)
mydict2 = nd.load('mydict')

"""
存储并读取一个从字符串映射到NDArray的字典
    {'x': 
    [ 1.  1.  1.]
    <NDArray 3 @cpu(0)>, 'y': 
    [ 0.  0.  0.  0.]
    <NDArray 4 @cpu(0)>}
"""
print(mydict2)


# 2、读写Gluon模型参数
class MLP(nn.Block):
    def __init__(self, **kwargs):
        super(MLP, self).__init__(**kwargs)
        self.hidden = nn.Dense(256, activation='relu')
        self.output = nn.Dense(10)

    def forward(self, x):
        return self.output(self.hidden(x))


net = MLP()
net.initialize()
X = nd.random.uniform(shape=(2, 20))
Y = net(X)

filename = 'mlp.params'
net.save_parameters(filename)

# 创建两个具有相同模型参数的实例
net2 = MLP()
net2.load_parameters(filename)

Y2 = net2(X)
"""
[[ 1.  1.  1.  1.  1.  1.  1.  1.  1.  1.]
 [ 1.  1.  1.  1.  1.  1.  1.  1.  1.  1.]]
<NDArray 2x10 @cpu(0)>
"""
print(Y2 == Y)
